# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

from experimental.prophet.graph import graph
from experimental.prophet.ops import timeseries
from experimental.prophet.ops import y_gen
from experimental.prophet.ops.control_flow import cond
from experimental.prophet.ops.constant import constant
from experimental.prophet.ops.util import to_duration_int


class FillQuery:
  def __init__(self,
               book,
               trade,
               latency,
               order_lifetime,
               maker_fee=None,
               taker_fee=None,
               timestamp=None):
    logging.warning('Deprecated')

    timestamp = timestamp or graph.default_timestamp()
    self._latency = to_duration_int(latency)
    self._order_lifetime = to_duration_int(order_lifetime)
    self._maker_fee = float(maker_fee) if maker_fee is not None else 0.
    self._taker_fee = float(taker_fee) if taker_fee is not None else 0.

    with graph.control_if(None):
      self._tacc_obj = graph.create_node('feature.TradeAccumulator',
                                         [timestamp, trade.price, trade.qty, trade.side], {
                                             'window_size_ns': self._order_lifetime
                                         }).outputs[0]
    self._book = book

  @property
  def latency(self):
    return self._latency

  @property
  def order_lifetime(self):
    return self._order_lifetime

  def query_buy(self, price, queue_pos=None):
    # t0: 0
    # t1: latency past
    # t2: order_lifetime past (order arrival time)
    # t3: latency + order_lifetime past (order submit time)
    with graph.control_if(None):
      if queue_pos is None:
        price_t1 = timeseries.time_shift(price, self.latency) if self.latency != 0 else price
        qpos_t0 = self._book.query_bid_qty(price_t1).level_qty
        qpos_t2 = timeseries.time_shift(qpos_t0, self.order_lifetime)
      else:
        qpos_t2 = graph.as_variable(queue_pos).cast(float)

      ask0_t2 = timeseries.time_shift(self._book.ask0.price, self.order_lifetime)
      price_t3 = timeseries.time_shift(price, self.latency + self.order_lifetime)

    filled = graph.create_node('feature.FillQuery', [self._tacc_obj, price_t3, qpos_t2], {
        'buy': True
    }).outputs[0]
    taker = price_t3 >= ask0_t2
    filled = cond(taker, True, filled)
    fill_price = cond(filled, cond(taker, ask0_t2, price_t3), 0.)
    fee = fill_price * cond(taker, self._taker_fee, self._maker_fee)

    with graph.control_if(None):
      return y_gen.time_shift_y([filled, taker, fill_price, fee],
                                self.latency + self.order_lifetime)

  def query_sell(self, price, queue_pos=None):
    # t0: 0
    # t1: latency past
    # t2: order_lifetime past (order arrival time)
    # t3: latency + order_lifetime past (order submit time)
    with graph.control_if(None):
      if queue_pos is None:
        price_t1 = timeseries.time_shift(price, self.latency) if self.latency != 0 else price
        qpos_t0 = self._book.query_ask_qty(price_t1).level_qty
        qpos_t2 = timeseries.time_shift(qpos_t0, self.order_lifetime)
      else:
        qpos_t2 = graph.as_variable(queue_pos).cast(float)

      bid0_t2 = timeseries.time_shift(self._book.bid0.price, self.order_lifetime)
      price_t3 = timeseries.time_shift(price, self.latency + self.order_lifetime)

    filled = graph.create_node('feature.FillQuery', [self._tacc_obj, price_t3, qpos_t2], {
        'buy': False
    }).outputs[0]
    taker = price_t3 <= bid0_t2
    filled = cond(taker, True, filled)
    fill_price = cond(filled, cond(taker, bid0_t2, price_t3), 0.)
    fee = fill_price * cond(taker, self._taker_fee, self._maker_fee)

    with graph.control_if(None):
      return y_gen.time_shift_y([filled, taker, fill_price, fee],
                                self.latency + self.order_lifetime)
